#!/usr/bin/env python
# -*- coding: utf-8 -*-

import random
from enum import IntEnum, auto
from asl_parser import *
import pickle

class DecodeError(Exception):
    pass

def extractBits(code, b0, bn):
    return (code>>b0) & ((1<<bn)-1)

def randInstGen(opcode, enc, bftypes):
    bsd = {'opcode':'[  0: 11] Opcode = ' + enc.mName + f'({opcode:#05x})'}

    code = opcode
    for bf, bfv in enc.mFields.items():
        if bf in {'optype', 'stype'}:
            continue

        bftypev = bftypes[bfv.mType]
        b0, bn = bfv.mRange[0]
        if bftypev.mVType == 'Enum':
            bv = random.choice(list(bftypev.mVK.keys()))

            if bftypev.mName in {'Reg', 'UReg'}:
                ralign = random.choices([0x3, 0x1, 0x0], weights=[6, 2, 2])[0]
                bv = bv & ~ralign
            bvs = bftypev.mVK[bv] + f'({bv:#x})'
        else:
            bv = random.randint(0, bftypev.mMaxV)
            if bftypev.mName == 'CMem':
                ralign = random.choices([0xf, 0x7, 0x3, 0x1, 0x0], weights=[4, 4, 4, 2, 2])[0]
                bv = bv & ~ralign
            bvs = f'{bv:#x}'
        code += bv<<bfv.mRange[0][0]

        bsd[bf] = f'[{b0:3d}:{b0+bn-1:3d}] {bfv.mType} {bf} = {bvs}'

    return code, bsd

def getBFAsm(enc, bf, bfv, bftypes):
    return bftypes[enc.mFields[bf].mType].mVK[bfv]

class RModiType(IntEnum):
    NONE = 0

    F_ABS = auto()    # F32
    F_NEG = auto()
    F_NEG_ABS = auto()

    # H_H0 = NONE
    H_H0_ABS = auto()    # FP16/BF16
    H_H0_NEG = auto()
    H_H0_NEG_ABS = auto()
    H_H1 = auto()
    H_H1_ABS = auto()
    H_H1_NEG = auto()
    H_H1_NEG_ABS = auto()

    # H_H1H0 = NONE
    H2_H1H0_ABS = auto()
    H2_H1H0_NEG = auto()
    H2_H1H0_NEG_ABS = auto()
    H2_H0H0 = auto()
    H2_H0H0_ABS = auto()
    H2_H0H0_NEG = auto()
    H2_H0H0_NEG_ABS = auto()
    H2_H1H1 = auto()
    H2_H1H1_ABS = auto()
    H2_H1H1_NEG = auto()
    H2_H1H1_NEG_ABS = auto()

    D_ABS = auto()    # F64
    D_NEG = auto()
    D_NEG_ABS = auto()

    I_NEG = auto()  # s32
    L_NEG = auto()  # s64 only for imad.wide?

    I_BITNOT = auto() # for all int?
    L_BITNOT = auto() # for all int?

    # BSEL_B0  // default = NONE?
    BSEL_B1 = auto()
    BSEL_B2 = auto()
    BSEL_B3 = auto()

    UNKNOWN = auto()
    END = auto()

class RModiProcessor:
    INVALID = 'INVALID'
    BSel = {'S0':'B0', 'S1':'B1', 'S2':'B2', 'S3':'B3'}
    HSel = {'S0':'H0', 'S1':'H1', 'S2':INVALID, 'S3':INVALID}

    VSrcDType = {
        'F2F':'srctype',
        'F2I':'ftype',
        'I2F':'itype',
        'FRND':'ftype',
        'F2F_64':'srctype',
        'F2I_64':'ftype',
        'I2F_64':'itype',
        'FRND_64':'ftype',
        'MUFU':'dtype',  # TODO: f64 only use upper 32bit
    }

    DMap = {'F32':'f', 'F64':'d', 'F16':'h', 'F64H': 'f', # 'F64H' means high 32bit for F64, only for MUFU
            'BF16':'h', 'U32':'i', 'S32':'i',
            'F16_V2':'h', 'BF16_V2':'h'}

    def __init__(self):
        # build rtype operand modi dict
        self.mRModis={
            (1, '') : RModiType.NONE,
            (2, '') : RModiType.NONE,
            (4, '') : RModiType.NONE,

            (1, 'ibitnot') : RModiType.I_BITNOT,
            (2, 'ibitnot') : RModiType.L_BITNOT,

            (1, 'ineg') : RModiType.I_NEG,
            (2, 'ineg') : RModiType.L_NEG,

            (1, 'fabs') : RModiType.F_ABS,
            (1, 'fneg') : RModiType.F_NEG,
            (1, 'fabs,fneg') : RModiType.F_NEG_ABS,

            (2, 'dabs') : RModiType.D_ABS,
            (2, 'dneg') : RModiType.D_NEG,
            (2, 'dabs,dneg') : RModiType.D_NEG_ABS,

            (1, 'H1_H0') : RModiType.NONE,
            (1, 'H1_H0,habs') : RModiType.H2_H1H0_ABS,
            (1, 'H1_H0,hneg') : RModiType.H2_H1H0_NEG,
            (1, 'H1_H0,habs,hneg') : RModiType.H2_H1H0_NEG_ABS,

            (1, 'H0_H0') : RModiType.H2_H0H0,
            (1, 'H0_H0,habs') : RModiType.H2_H0H0_ABS,
            (1, 'H0_H0,hneg') : RModiType.H2_H0H0_NEG,
            (1, 'H0_H0,habs,hneg') : RModiType.H2_H0H0_NEG_ABS,

            (1, 'H1_H1') : RModiType.H2_H1H1,
            (1, 'H1_H1,habs') : RModiType.H2_H1H1_ABS,
            (1, 'H1_H1,hneg') : RModiType.H2_H1H1_NEG,
            (1, 'H1_H1,habs,hneg') : RModiType.H2_H1H1_NEG_ABS,

            # Only for HSel2.F32
            (1, 'F32') : RModiType.NONE,
            (1, 'F32,habs') : RModiType.F_ABS,
            (1, 'F32,hneg') : RModiType.F_NEG,
            (1, 'F32,habs,hneg') : RModiType.F_NEG_ABS,

            (1, 'B0') : RModiType.NONE,
            (1, 'B1') : RModiType.BSEL_B1,
            (1, 'B2') : RModiType.BSEL_B2,
            (1, 'B3') : RModiType.BSEL_B3,

            (1, 'H0') : RModiType.NONE,
            (1, 'H0,habs') : RModiType.H_H0_ABS,
            (1, 'H0,hneg') : RModiType.H_H0_NEG,
            (1, 'H0,habs,hneg') : RModiType.H_H0_NEG_ABS,

            # For F2F/F2I, only H0 is allowed for F32/F64
            (1, 'H0') : RModiType.NONE,
            (1, 'H0,fabs') : RModiType.F_ABS,
            (1, 'H0,fneg') : RModiType.F_NEG,
            (1, 'H0,fabs,fneg') : RModiType.F_NEG_ABS,

            (2, 'H0') : RModiType.NONE,
            (2, 'H0,dabs') : RModiType.D_ABS,
            (2, 'H0,dneg') : RModiType.D_NEG,
            (2, 'H0,dabs,dneg') : RModiType.D_NEG_ABS,

            (1, 'H1') : RModiType.H_H1,
            (1, 'H1,habs') : RModiType.H_H1_ABS,
            (1, 'H1,hneg') : RModiType.H_H1_NEG,
            (1, 'H1,habs,hneg') : RModiType.H_H1_NEG_ABS,

            (1, 'buf') : RModiType.NONE,
            (2, 'buf') : RModiType.NONE,
            (4, 'buf') : RModiType.NONE,
            (1, 'nobuf') : RModiType.NONE,
            (2, 'nobuf') : RModiType.NONE,
            (4, 'nobuf') : RModiType.NONE,
        }

        # build modifier dependent modifier dict
        self.mDepModis = {}
        vsel = {}
        for v,r in self.BSel.items():
            for t in ['S8','U8']:
                vsel[(v,t)] = r
        for v,r in self.HSel.items():
            for t in ['S16', 'U16']:
                vsel[(v,t)] = r

        self.mDepModis['CvtVSel'] = vsel

        vpsel = {}
        for v,r in self.BSel.items():
            for t in ['E2M1']:
                vpsel[(v,t)] = r
        for v,r in self.HSel.items():
            for t in ['E5M2','E4M3','E8','E3M2','E2M3']:
                vpsel[(v,t)] = r
        self.mDepModis['CvtVPSel'] = vpsel

        self.mDepModis['CvtINegX'] = {('False','NoX') : None,
                                  ('False','X'  ) : None,
                                  ('True' ,'NoX') : 'neg',
                                  ('True' ,'X'  ) : 'bitnot'}

    def trDepModi(self, vval, mval, func):
        if func in self.mDepModis:
            r = self.mDepModis[func]
            if (vval, mval) in r:
                return r[(vval, mval)]
            else:
                return None
        elif func == 'CvtFImm':
            return None
        else:
            raise DecodeError('Unknown modi func {func}!!!')

    def getRModiType(self, rmodis, rlen, bfvals, bftypes, enc, operandInfo):
        optype = bftypes['Optype'].mVK[bfvals['optype']]
        s_dtype = RModiProcessor.getSrcDType(bfvals, bftypes, enc)

        rmoditag = []
        for bf, abf, atype, bv in rmodis:

            if 'AsmFormat' in operandInfo.mPData and bf in operandInfo.mPData['AsmFormat']:
                func, var, modi = operandInfo.mPData['AsmFormat'][bf]

                vval = getBFAsm(enc, bf, bv, bftypes)
                mval = getBFAsm(enc, modi, bfvals[modi], bftypes)
                tr_modi = self.trDepModi(vval, mval, func)
                if tr_modi is None:
                    continue
                else:
                    if tr_modi in {'neg', 'bitnot'}:
                        atype = tr_modi
                    else:
                        rmoditag.append(tr_modi)
                        continue

            # SignModi should associated with dtype
            if atype in {'neg', 'abs'}:
                if bv != 0: # TODO:???
                    rmoditag.append(s_dtype + atype)
            elif atype == 'bitnot':
                if bv != 0: # TODO:???
                    rmoditag.append('i' + atype)
            elif atype == 'not':# not only for P/UP
                if bv != 0: # TODO:???
                    rmoditag.append(atype)
            else:
                mtag = bftypes[enc.mFields[bf].mType].mVK[bv]
                rmoditag.append(mtag)

        rmoditag.sort()
        rmkey = (rlen, ','.join(rmoditag))
        if rmkey in self.mRModis:
            return rmoditag, self.mRModis[rmkey]
        else:
            return rmoditag, RModiType.UNKNOWN

    @staticmethod
    def getSrcDType(bfvals, bftypes, enc):
        ''' Get source data type:

            Only for neg/abs/bitnot
        '''

        optype = bftypes['Optype'].mVK[bfvals['optype']]
        if optype in RModiProcessor.VSrcDType:
            bf_src = RModiProcessor.VSrcDType[optype]
            bfval = bfvals[bf_src]
            bftype = enc.mFields[bf_src].mType
            srctype = bftypes[bftype].mVK[bfval]
            if srctype in RModiProcessor.DMap:
                # if optype=='MUFU' and srctype=='F64':  # MUFU.F64 is renamed to F64H, thus no need for special treatment here.
                #     return 'f'
                # else:

                return RModiProcessor.DMap[srctype]
            else:
                return 'i'
        else:
            if optype[0] == 'H':
                return 'h'
            elif optype[0] == 'D':
                return 'd'
            elif optype[0] == 'F':
                return 'f'
            else:
                return 'i'  # TODO: ???

class PreDecoder:
    INSTRUCTION_SIZE = 16
    OPCODE_MASK = 0xfff

    FUNC_IMM_BFS = {'ULOP3.lut',
                    'LOP3.lut',
                    'UPLOP3.lut',
                    'PLOP3.lut',
                    'LEA.shiftamt',
                    'ULEA.shiftamt',
                    'HMMA_SP.mdsel',
                    'IMMA_SP.mdsel',
                    'QMMA_SP.mdsel'}

    IMPLICIT_OPERAND = {'PR', 'UPR'}

    CONTROL_CODE_BFS = {'ReqSB', 'WrSB', 'Sched', 'Trap'}

    GPR_BFS = {'Reg', 'UReg', 'Pred', 'UPred'}

    def __init__(self, repos:dict, funcs:dict, bftypes:dict):
        opc = {}
        for k, opv in repos.items():
            if opv.mType == 'Opcode':
                enc = opv.mData['__Encoding']
                opcode = enc.mFieldConstants['optype'] + (enc.mFieldConstants['stype']<<8)
                # print(f"Opcode {opv.mName} = {opcode:#05x}")
                opc[opcode] = opv

        self.mOpcodeRepos = opc
        self.mOpFuncs = funcs
        self.mBFTypes = bftypes

    def decode(self, code:int, *, return_exception=False, fout=None):
        ''' decode instruction code.

            TODO: set Func to 0xffffffff for error encoding?
                  or just 0b11??????...
        '''

        wlog = lambda s: '' if fout is None else fout.write(s)

        if code<0 or (code>>(self.INSTRUCTION_SIZE*8))>0:
            raise DecodeError(f"Instruction code out of range {code:#x}!")

        # get opcode
        opcode = code & self.OPCODE_MASK

        if opcode not in self.mOpcodeRepos:
            raise DecodeError(f"Unknown opcode {opcode:#05x} in {code:#018x}!")

        enc = self.mOpcodeRepos[opcode].mData['__Encoding']
        opinfo = self.mOpcodeRepos[opcode].mData['__OperandInfo']
        excep = self.mOpcodeRepos[opcode].mData['__Exception']

        # collect all operand list
        opr_list = []
        for op in opinfo.mPData['Order']:
            if isinstance(op, list):
                opr_list.extend(op[1:])
            elif op in {'PR', 'UPR'}:
                pass
            else:
                opr_list.append(op)

        # check unused bits
        optype = enc.mFields['optype'].mVal.strip(' =')
        bfval_dict = {}

        d = {'func':[], 'cc':[], 'opr':[], 'opr_modi':{}, 'extra':[]}

        # enumerate bitfields and categorize
        for bf, bfv in enc.mFields.items():
            bftype = bfv.mType
            b0, bn = bfv.mRange[0]  # TODO: currently all bitfields have ONE range
            bv = extractBits(code, b0, bn)

            bfval_dict[bf] = bv

            if bftype == 'Optype':                # optype goes to .func
                d['func'].append((bf, bfv, bv))
            elif bftype in {'SType', 'USType'}:   # ignore stype ???
                continue
            elif bftype in self.GPR_BFS:          # io Reg/UReg/Pred/UPred
                d['opr'].append((bf, bfv, bv))
            elif bftype == 'BReg':                # BReg treated as UReg, treatments depend on optype
                d['opr'].append((bf, bfv, bv))    # NOTE: breg B# will be read again in func
                                                  #       thus all ops to R# will be discarded.
                                                  # TODO: BMOV_R_B/BMOV_B_R ????
            elif '.' in bf:                       # operand associated var
                if bftype in {'ADDRMode', 'RStride'}:        #
                    d['func'].append((bf, bfv, bv))
                else:
                    # d['opr_a'].append(bfv)
                    abf, atype = tuple(bf.split('.'))   # TODO: modifier dependent op?
                    if abf in d['opr_modi']:
                        d['opr_modi'][abf].append((bf, abf, atype, bv))
                    else:
                        d['opr_modi'][abf] = [(bf, abf, atype, bv)]
            elif bftype in self.CONTROL_CODE_BFS: # control codes go to .cc
                d['cc'].append((bf, bfv, bv))
            elif f'{optype}.{bf}' in self.FUNC_IMM_BFS:
                d['func'].append((bf, bfv, bv))
            elif bf in opr_list:                  # other operands, usually imm
                if bn>32:                         # Wide bitfields go to extra(bn>32, only for Imm)
                    d['extra'].append((bf, bfv, bv))
                else:
                    d['opr'].append((bf, bfv, bv))
            else: # other modifiers go to .func
                d['func'].append((bf, bfv, bv))

        flag, msg = enc.checkEncodingError(bfval_dict, opinfo, excep, self.mBFTypes, maxreg=254)
        if not flag :
            raise DecodeError('InvalidEncoding! ' + msg)

        # eval reg bitwidth
        bw_dict = opinfo.evalBitwidth(bfval_dict)
        rlen_dict = {}
        for bf, bw in bw_dict.items():
            if enc.mFields[bf].mType in {'Reg', 'UReg', 'CMem', 'BReg'} :
                rl = (bw + 31) // 32
                if rl not in {1, 2, 4}:
                    raise DecodeError(f'Invalid reg len {rl} for {bf}!')
                rlen_dict[bf] = rl

        # calc func key
        func_key = 0
        for _, bfv, bv in d['func']:
            func_key += bv<<bfv.mRange[0][0]

        if optype in OPTYPES_FUNC_WITH_STYPE:
            func_key += enc.mFieldConstants['stype']<<8

        if func_key not in self.mOpFuncs:
            raise DecodeError(f'Unknown func key {func_key:#x}!')
        func_tag = self.mOpFuncs[func_key]

        opr_iset = set(opinfo.mPData['InList'])
        opr_oset = set(opinfo.mPData['OutList'])

        # all Reg/UReg/CMem/Pred/UPred + IMM_Not_In_Func

        opr_io = {'pmask':0, 'rmask':0, 'pi':[], 'po':[], 'ri':[], 'ro':[]}
        for bf in opr_list:
            if f'{optype}.{bf}' in self.FUNC_IMM_BFS:
               continue

            bftype = enc.mFields[bf].mType

            if bftype in {'Pred', 'UPred'}:
                if bf in opr_iset:
                    opr_io['pi'].append((bf, bftype, bfval_dict[bf]))
                else:
                    opr_io['po'].append((bf, bftype, bfval_dict[bf]))
            elif bftype in {'Reg', 'UReg'}:
                if bf in opr_iset:
                    opr_io['ri'].append((bf, bftype, bfval_dict[bf], rlen_dict[bf]))
                else:
                    opr_io['ro'].append((bf, bftype, bfval_dict[bf], rlen_dict[bf]))
            elif bftype == 'BReg':
                opr_io['ri'].append((bf, bftype, bfval_dict[bf], rlen_dict[bf]))
            else:
                opr_io['ri'].append((bf, bftype, bfval_dict[bf], 1)) # Imm longer than 32bit will go to extra

        wlog('Func:\n')
        wlog(f'  func_key = {func_key:#034x}\n')
        wlog(f'  func_tag = {func_tag}\n')
        for bf, bfv, bv in d['func']:
            wlog(f'  {bfv.mType} {bf} = {bv:#x}\n')

        rmp = RModiProcessor()
        def w_opr_modi(bf, bftype):
            if bf in d['opr_modi'] and bftype in {'Reg','UReg'}:
                for (mbf, _, _, bv) in d['opr_modi'][bf]:
                    wlog(f'    {" "*len(bftype)} {mbf} = {bv:#x}\n')
                rmoditag = rmp.getRModiType(d['opr_modi'][bf], rlen_dict[bf], bfval_dict, self.mBFTypes, enc, opinfo)
                wlog(f'    rmodi = {rmoditag}\n')

        wlog('Operands:\n')
        wlog('  p_in:\n')
        for bf, bftype, bv in opr_io['pi']:
            wlog(f'    {bftype} {bf} = {bv:#x}\n')
            w_opr_modi(bf, bftype)
        wlog('  p_out:\n')
        for bf, bftype, bv in opr_io['po']:
            wlog(f'    {bftype} {bf} = {bv:#x}\n')
            w_opr_modi(bf, bftype)
        wlog('  r_in:\n')
        for bf, bftype, bv, rl in opr_io['ri']:
            wlog(f'    {bftype} {bf} = {bv:#x}, rlen={rl}\n')
            w_opr_modi(bf, bftype)
        wlog('  r_out:\n')
        for bf, bftype, bv, rl in opr_io['ro']:
            wlog(f'    {bftype} {bf} = {bv:#x}, rlen={rl}\n')
            w_opr_modi(bf, bftype)

        wlog('Extra:\n')
        for bf, bfv, bv in d['extra']:
            wlog(f'  {bf}={bv:#x}\n')

    def decodeBytes(self, code_bytes:str):
        pass

def listAllOperandModi(opgroups, fout=None):
    do_fclose = False
    if fout is not None and isinstance(fout, str):
        fout = open(fout, 'w+')
        do_fclose = True

    repos = {}
    for op, opv in opgroups.items():
        if opv.mType != 'Opcode':
            continue

        enc = opv.mData['__Encoding']

        d = {}
        for bf, bfv in enc.mFields.items():
            if '.' in bf:
                vbf, vop = tuple(bf.split('.'))
                vop = bfv.mType+'.'+vop
                if vbf in d:
                    d[vbf].append(vop)
                else:
                    d[vbf] = [vop]

        for vbf, vops in d.items():
            vops.sort()
            vop_tag = '[' + (', '.join(vops)) +']'
            tag = f'{op}.{vbf}'
            if vop_tag not in repos:
                repos[vop_tag] = [tag]
            else:
                repos[vop_tag].append(tag)

    def wlog(s):
        print(s)
        if fout is not None:
            fout.write(s+'\n')

    for vops, tags in repos.items():
        wlog('-'*80)
        wlog(f'{vops} :')
        for tag in tags:
            wlog(f'  {tag}')

    if do_fclose:
        fout.close()

def testPreDecoder(asl_data, logfile='predecode.log'):
    pd = PreDecoder(asl_data['mOpGroups'], asl_data['mOpFuncs'], asl_data['mTypes'])

    # ks = {0x902:'FFMA_RRR'}
    opcodes = pd.mOpcodeRepos.keys()

    total_try_count = 0
    total_pass_count = 0

    NMaxTry = 128
    # NMinPass = 4
    NMaxPass = 16

    dec_info_dict = {}
    with open(logfile, 'w+') as flog:
        for opcode in opcodes:
            repos = pd.mOpcodeRepos[opcode]
            itry = 0
            cnt = 0
            while cnt<NMaxPass and itry<NMaxTry:
                itry += 1
                flog.write('-'*80 + '\n')
                flog.write(f'# Opcode = {repos.mName} ({opcode:#05x})\n')

                code, bsd = randInstGen(opcode, repos.mData['__Encoding'], asl_data['mTypes'])
                flog.write('InstructionCode:\n')
                for k, v in bsd.items():
                    flog.write('    ' + v + '\n')
                try:
                    pd.decode(code, fout=flog)
                    cnt += 1
                except DecodeError as de:
                    flog.write(f'!!! DecodeError! {de.args}\n')

            dec_info_dict[opcode] = itry, cnt
            total_try_count += itry
            total_pass_count += cnt

        flog.write('='*120 + '\n')
        flog.write('Summary\n')
        flog.write('='*120 + '\n')
        flog.write(f'  TotalTryCount = {total_try_count}\n')
        flog.write(f'  TotalPassCount = {total_pass_count}\n')
        flog.write(f'  Try/Pass count:\n')
        for i, (k, v) in enumerate(dec_info_dict.items()):
            opname = pd.mOpcodeRepos[k].mName
            if v[1] == 0:
                flog.write(f'    {i:3d}  {opname:<20s} : ({v[0]:4d}, {v[1]:4d}) !!!!\n')
            else:
                flog.write(f'    {i:3d}  {opname:<20s} : ({v[0]:4d}, {v[1]:4d})\n')

if __name__ == '__main__':

    fname = r'D:\Notes\notes\docs\ISASpec\build\asl_data.pickle'

    with open(fname, 'rb') as fin:
        asl_data = pickle.load(fin)

    # listAllOperandModi(asl_data['mOpGroups'], 'opr_modis.log')
    testPreDecoder(asl_data)

